package plan

import (
	"fmt"
	"mini-sql/expression"
	"mini-sql/model"
	"mini-sql/storage"
)

type UpdatePlanNode struct {
	source          PlanNode
	updateStatement *model.UpdateStatement
	storePool       *storage.StorePool
	tableStore      storage.StoreInterface
}

func NewUpdatePlanNode(source PlanNode, updateStatement *model.UpdateStatement, storePool *storage.StorePool) PlanNode {
	return &UpdatePlanNode{
		source:          source,
		updateStatement: updateStatement,
		storePool:       storePool,
	}
}

func (node *UpdatePlanNode) Prepare() error {
	tableName, err := expression.GetTableName(node.updateStatement.TableRef)
	if err != nil {
		return err
	}
	if tableName == "" {
		return fmt.Errorf("SelectPlan: tableName is empty")
	}
	tableStore, err := node.storePool.GetTableStore(tableName)
	if err != nil {
		return err
	}
	node.tableStore = tableStore

	return node.source.Prepare()
}

func (node *UpdatePlanNode) GetNextTuple() (*model.TableTuple, error) {
	tuple, err := node.source.GetNextTuple()
	if err != nil {
		return nil, err
	}
	if tuple == nil {
		return nil, nil
	}
	// 执行更新
	for _, assignExpr := range node.updateStatement.AssignList {
		fieldIndex := tuple.GetValueIndexByName(assignExpr.Name)
		newValue := expression.EvalExpr(assignExpr.Value, tuple).ToValue()
		if fieldIndex >= 0 {
			tuple.Fields[fieldIndex].Value = newValue
		} else {
			tuple.Fields = append(tuple.Fields, model.TableField{
				Name:  assignExpr.Name,
				Value: newValue,
			})
		}
	}
	err = node.tableStore.UpdateRow(*tuple)
	if err != nil {
		return nil, err
	}
	return tuple, nil
}
